#!/usr/bin/env python

from BeautifulSoup import BeautifulSoup
from Bio.PDB import *
from optparse import OptionParser

import urllib, csv, re, sys
import textwrap, urllib2, os
import xml.etree.ElementTree as ET
from Bio import ExPASy
import MySQLdb
import dbtbl



# This script will give the protein sequences of PDZ domains that bind to the
# peptides determined by peptide_extract.pl by searching the domino file and
# uniprot database. bind_from and bind_to from domino file and full sequence
# from database.

# 3 things needed to run this script
# 1. peptide1.out file generated by peptide1.py
# 2. domino database mitab file can download from DOMINO database FTP
# 3. Uniprot database (.xml) file in uniprot/ folder

'''
For PDZ <--> peptides
load peptide2 file, using dict. DOMINOID as key, peptide as value
Protein A + proteinB == '-'
link by DOMINO ID
#------------------------------------#
PDZ <--> other proteins
1. find which is the PDZ, A or B
2. find what are the binding sites
     PDZ domain --> more than ONE? overlaps?
     third protein --> binding sites at c-term? almost c-term?
     or somewhere else
3. output interactions

Note:
A-B interactions will have redundences. i.e. overlapping PDZ domains
There will be some cases that B not interacts at the c-term
and very likely cases that interacts near c-terminal.
How to distinguish near c-term and in the middle???

'''

ns = '{http://uniprot.org/uniprot}'
#res = urllib2.urlopen('http://www.uniprot.org/uniprot/Q9NZW5.xml')

def GetName(self):
    '''Get the name of the protein'''
    try:
        name = (protein.find(ns+'protein')
            .find(ns+'recommendedName')
                .findtext(ns+'fullName'))
    except:
        name = (protein.find(ns+'protein')
                .find(ns+'submittedName')
                .findtext(ns+'fullName'))
    return name
    
def GetAccessions(self):
    '''Get all accession codes for this protein'''
    accession_list = []
    accessions = protein.findall(ns+'accession')
    for accession in accessions:
        accession_list.append(accession.text)
    return accession_list 

def GetGeneName(self):
    '''Get the gene name'''
    try:
        gene_name = protein.find(ns+'gene').findtext(ns+'name')
    except:
        # Have to invent short names for proteins that don't have one
        try:
            fullname = (protein.find(ns+'protein')
                        .find(ns+'recommendedName')
                        .findtext(ns+'fullName'))
        except:
            fullname = (protein.find(ns+'protein')
                        .find(ns+'submittedName')
                        .findtext(ns+'fullName'))
        gene_name = fullname[0:13]
    return gene_name

def GetSequence(self):
    '''Retrieve the protein sequence'''
    sequence = protein.findtext(ns+'sequence')
    seq_clean = re.sub("\s+", "", sequence)
    #seq_wrapped = textwrap.fill(seq_clean, 60)
    return seq_clean

def GetOrganism(self): 
    '''Retrieve the protein organism'''
    organism_name = 'None'
    organism_list = protein.find(ns+'organism').findall(ns+'name')
    for organism in organism_list:  
        if organism.get('type') == 'scientific':
            organism_name = organism.text
    return organism_name

def GetLocation(self):
    '''return length with start and stop'''
    
    l_begin = (protein.find(ns+'feature').
               find(ns+'location').
               find(ns+'begin').
               get('position'))
    l_end = (protein.find(ns+'feature')
             .find(ns+'location')
             .find(ns+'end')
             .get('position'))
    return int(l_begin), int(l_end)
    

def GetDomain(self):
    '''Return domains with start and stop positions'''
    domain_dict = {}
    domain_list = []
    domains = protein.findall(ns+'feature')
    # This loop seems cumbersome
    for feature in domains:
        if feature.get('type') == 'domain':
            domain_list.append(feature)
            for domain in domain_list:
                domain_start = (feature.find(ns+'location')
                                .find(ns+'begin')
                                .get('position'))
                domain_stop = (feature.find(ns+'location')
                               .find(ns+'end')
                               .get('position'))
                domain_span = domain_start+'-'+domain_stop
                domain_dict[feature.get('description')]= domain_span
    return domain_dict

def GetLength(self):
    '''return the sequence length of protein'''
    length = protein.find(ns+'sequence').get('length')
    return int(length)
    
def uniprot_ext(proteinID, BindingSiteS, BindingSiteE):
    #
    #return proteinID, BindingSiteS, BindingSiteE, sequence
    #----------------------------------
    
    nameA = '-'
    lenA = '-'
    accessionsA = '-'
    genenameA = '-'
    organismA = '-'
    domain_dictA = '-'
    sequenceA = '-'
    #------------------------------------
    #if os.path.getsize(
    #    '/home/sdai/workspace/PDZ/uniprot/'+proteinID+'.xml') != 0:
    if os.path.getsize('uniprot/'+proteinID+'.xml') != 0:
        #=0 means this entry has already deleted by uniprot
        #tree = ET.parse('/home/sdai/workspace/PDZ/uniprot/'+proteinID+'.xml')
        tree = ET.parse('uniprot/'+proteinID+'.xml')
        for item in tree.getiterator(ns+'uniprot'):
            global protein
            protein = item.findall(ns+'entry')[0]
            nameA = GetName(item)
            lenA = GetLength(item)
            accessionsA = GetAccessions(item)
            genenameA = GetGeneName(item)
            organismA = GetOrganism(item)
            domain_dictA = GetDomain(item)
            sequenceA = GetSequence(item)
        '''
        if abs(BindingSiteS-BindingSiteE) > 100 or (BindingSiteS == 0 and
                                                    (BindingSiteE > 100
                                                     or BindingSiteE ==0)):
                
                print proteinID, BindingSiteS, BindingSiteE
        '''
        
        if BindingSiteS == 0 and BindingSiteE == 0:
            return (proteinID, BindingSiteS, BindingSiteE, sequenceA)
            
        elif BindingSiteE == 0:
            return (proteinID, BindingSiteS, BindingSiteE,
                    sequenceA[BindingSiteS-1:])
        elif BindingSiteS ==0:
            return (proteinID, BindingSiteS, BindingSiteE,
                    sequenceA[:BindingSiteE-1])
        elif BindingSiteE-BindingSiteS <= 5:
            return (proteinID, BindingSiteS, BindingSiteE,
                    sequenceA[BindingSiteE-6:BindingSiteE-1])
        else:
            return (proteinID, BindingSiteS, BindingSiteE,
                    sequenceA[BindingSiteS-1:BindingSiteE-1])
    else:
        return

def load_peptide(filename):
    # load peptide1.out generated by peptide1.py
    pepdb = dict()
    input_table = csv.reader(open(filename), delimiter=' ')
    for row in input_table:
        pepdb[row[1]] = (row[0], row[2])
    return pepdb

def masking(filename, uni_acc, peptide):
    #lOAD Validation dataset
    validation_table = csv.reader(open(filename))
    for row in validation_table:
        uniprot_acc = row[0]
        vpeptide = row[3]
        if uniprot_acc == uni_acc and peptide == vpeptide[-5:]:
            return True
    return False
    
#-------------------------------------------------------------------##



###================================================================###
###----------------------------------------------------------------###
###================================================================###
#pep = load_peptide('/home/sdai/workspace/PDZ/peptide1.out')

def main():
    
    pep = load_peptide('peptide2.out')
    #domino_table = csv.reader(open(
    #        '/home/sdai/workspace/PDZ/2009-10-22-domino-pdz-binary.mitab26'),
    #        delimiter='\t')
    domino_table = csv.reader(open('2009-10-22-domino-pdz-binary.mitab26'),
                              delimiter='\t')
    #csvreader.
    uni_proteinID = set()
    proteinIDA = list()
    proteinIDB = list()
    BindingSiteA = list()
    BindingSiteB = list()
    BindingSiteAe = list()
    BindingSiteBe = list()
    BindingSiteAs = list()
    BindingSiteBs = list()
    BindingDomainA = list()
    BindingDomainB =list()
    proteinIDAdb = list()
    proteinIDBdb = list()
    DominoID = list()
    taxidA = list()
    taxidB = list()
    wholeSeq = list()
    
    
    totalall = 0
    for row in domino_table:
        #totalall += 1
        #iiden = row[13] #Interaction identifier(s)
        uni_proteinID.add(row[0].split(':')[1]) #interactor ID A
        
        #uni_proteinIDB.add(row[0].split(':')[1])
        #load binding pocket
        
        try:
            proteinIDA.append(row[0].split(':')[1])
            proteinIDAdb.append(row[0].split(':')[0])
        except (IndexError):
            proteinIDA.append('-')
            proteinIDAdb.append('-')
            
        try:
            proteinIDB.append(row[1].split(':')[1])
            proteinIDBdb.append(row[1].split(':')[0])
        except (IndexError):
            proteinIDB.append('-')
            proteinIDBdb.append('-')
            
            
        try:
            BindingDomainA.append(row[32].split(';')[2])
        except (IndexError):
            BindingDomainA.append('-')
            
        try:
            BindingDomainB.append(row[33].split(';')[2])
        except (IndexError):
            BindingDomainB.append('-')
            
        try: 
            BindingSiteA.append(row[32].split(';')[0])
            
        except (IndexError):
            BindingSiteA.append('-')
    
        try:
            BindingSiteAe.append(int(row[32].split(';')[0].split('-')[1]))
        except (ValueError, IndexError):
            BindingSiteAe.append(0)
            
        try:
            BindingSiteAs.append(int(row[32].split(';')[0].split('-')[0]))
        except (ValueError, IndexError):
            BindingSiteAs.append(0)
            
        try:
            BindingSiteB.append(row[33].split(';')[0])
        except (ValueError, IndexError):
            BindingSiteB.append('-')
    
        try:
            BindingSiteBe.append(int(row[33].split(';')[0].split('-')[1]))
        except (ValueError, IndexError):
            BindingSiteBe.append(0)
    
        try:
            BindingSiteBs.append(int(row[33].split(';')[0].split('-')[0]))
        except (ValueError, IndexError):
            BindingSiteBs.append(0)
            
        try: 
            taxidA.append(row[9].split(':')[1].split('(')[0])
        except:
            taxidA.append('-')
            
        try:
            taxidB.append(row[10].split(':')[1].split('(')[0])
        except:
            taxidB.append('-')
            
    #    Get full protein sequence from Uniprot by AC  - should be easy.
    #    try:
    #        wholeSeq.append(row)
                    
        try:
            DominoID.append(row[13].split(':')[1])
        except:
            DominoID.append('-')
    #----------------------#
    
    # 1. load the binding pocket pocket
    
    # 2. extract xx;xx;xx
    
    # 3. extract size from Uniprot (uniprot parser)
            
    # 4. compare size  point to range .....
    
    # 
    #----------------------#
    #interactorAc
    #print len(BindingSiteB)
    #print len(BindingSiteA)
    #print len(proteinIDA)
    print "total:", totalall, len(BindingSiteA)
    pdzdomains = dict()
    err_count = 0
    for iter in range(len(BindingSiteA)):
    
        #Protein A + Protein B both PDZ
        #Protein A is, protein B is not
        #Protein A is not, protein B is
        #Protein A + peptide
        
        if (proteinIDAdb[iter] == 'uniprotkb' and
            (re.match('interpro:IPR001478', BindingDomainA[iter]))):
            #Protein A is a PDZ domain
            #print proteinIDA[iter], BindingSiteAs[iter], BindingSiteAe[iter]
            
            try:
                (protA, Astart, Aend, Aseq) = uniprot_ext(proteinIDA[iter],
                                                          BindingSiteAs[iter],
                                                          BindingSiteAe[iter])
            except Exception, e:
                print "Error:", e, proteinIDA[iter]
                err_count += 1
            
            if (proteinIDBdb[iter] == 'uniprotkb'
                and proteinIDB[iter] != '-'
                and (re.match('interpro:IPR001478', BindingDomainB[iter]))):
                #there are only ONE instant of protein A and B are both PDZs.
                #Q02410 autoinhibition of X11/alpha
                
                try:
                    (protB, Bstart, Bend, Bseq) = uniprot_ext(proteinIDB[iter],
                                                              BindingSiteBs[iter],
                                                              BindingSiteBe[iter])
                except Exception, e:
                    print "Error:", e, proteinIDB[iter]
                    err_count += 1
                    
                key = (proteinIDA[iter], proteinIDB[iter])
                if key in pdzdomains:
                    (ks,ke,ksa,ksb) = pdzdomains[key]
                    if (abs(ke-ks) > abs(Aend - Astart)) or (ke == ks):
                        pdzdomains[key] = (Astart, Aend, Aseq, Bseq[-5:])
                else:
                    pdzdomains[key] = (Astart, Aend, Aseq, Bseq[-5:])
                    
                    
            elif proteinIDB[iter] == '-':
                #protein A - peptide (ID_B empty)
                #record protein A in pdzdomains
                #pep[Domino ID] = (acc, peptide)
                
                if DominoID[iter] in pep:
                    
                    key = (pep[DominoID[iter]][0],pep[DominoID[iter]][1][-5:])
                    if key in pdzdomains:
                        (ks,ke,ksa,ksb) = pdzdomains[key]
                        if (abs(ke-ks) > abs(Aend - Astart)) or (ke == ks):
                            pdzdomains[key] = (Astart, Aend, Aseq, key[1])
                    else:
                        pdzdomains[key] = (Astart, Aend, Aseq, key[1])
                        #print proteinIDA[iter] == pep[DominoID[iter]][0]
                else:
                    print DominoID[iter]
                    print proteinIDA[iter]
                    totalall += 1
                    #raw_input('wait')
    
    
            else:
                #protein A<PDZ, proteinB<not PDZ
                #record protein A in pdzdomains
                #print 'A is B not'
                try:
                    (protB, Bstart, Bend, Bseq) = uniprot_ext(proteinIDB[iter],
                                                              BindingSiteBs[iter],
                                                              BindingSiteBe[iter])
                except(TypeError):
                    print proteinIDB[iter], ' has been deleted.'
                    err_count += 1
    
                key = (proteinIDA[iter], proteinIDB[iter])
                if key in pdzdomains:
                    (ks,ke,ksa,ksb) = pdzdomains[key]
                    if (abs(ke-ks) > abs(Aend - Astart)) or (ke == ks):
                        pdzdomains[key] = (Astart, Aend, Aseq, Bseq[-5:])
                else:
                    pdzdomains[key] = (Astart, Aend, Aseq, Bseq[-5:])
                    
        elif (proteinIDAdb[iter] == 'uniprotkb'
              and proteinIDBdb[iter] == 'uniprotikb'
              and (re.match('interpro:IPR001478', BindingDomainB[iter]))):
            #protein A is not PDZ, protein B is PDZ
            (protA, Astart, Aend, Aseq) = uniprot_ext(proteinIDA[iter],
                                                      BindingSiteAs[iter],
                                                      BindingSiteAe[iter])
            (protB, Bstart, Bend, Bseq) = uniprot_ext(proteinIDB[iter],
                                                      BindingSiteBs[iter],
                                                      BindingSiteBe[iter])
            #key will have the pdz domain in front
            key = (proteinIDB[iter], proteinIDA[iter])
            if key in pdzdomains:
                (ks,ke,ksa,ksb) = pdzdomains[key]
                if (abs(ke-ks) > abs(Bend - Bstart)) or (ke == ks):
                    pdzdomains[key] = (Bstart, Bend, seqB, Aseq[-5:])
            else:
                pdzdomains[key] = (Bstart, Bend, seqB, Aseq[-5:])
            
            
    
    # download uniprot filen
    #locale database
    #for iter in uni_proteinID:
    #    uniprot_file = urllib.urlretrieve
    #        ('http://www.uniprot.org/uniprot/'+iter+'.xml', iter+'.xml')
    
    #10 fold cross-validation database
    #10 databases: pdzvalid_1 to pdzvalid_10
    #using different mask files
    dbuser = "root"
    dbpassword = "198544"
    #for database_num in range(1, 11):
        
    #dbname = "pdzvalid_"+str(database_num)

    tbl_sdr = set()
    tbl_sdrdom = set()
    tbl_pdzdom = dict()
    tbl_peptide = set()
    tbl_pepdom = set()
    #vfile = 'test_group'+str(database_num-1)+'.dat'
    total_count = 0
    for i in pdzdomains:
        #(tbl_sdr, tbl_sdrdom, tbl_pdzdom, tbl_peptide, tbl_pepdom)
        #if masking(vfile ,i[0], pdzdomains[i][3]):
        #    print "THIS ENTRY IS MASKED."
        #else:        
        hold = dbtbl.build(
            i[0], i[1], pdzdomains[i][0], pdzdomains[i][1], pdzdomains[i][2],
            pdzdomains[i][3], tbl_sdr, tbl_sdrdom, tbl_pdzdom, tbl_peptide,
            tbl_pepdom)
        #    #dbtbl.build(protein A, protein B, start position, end position,
            #           sequence, peptide, table holders)
        if hold:
            total_count += 1
            (tbl_sdr, tbl_sdrdom, tbl_pdzdom, tbl_peptide, tbl_pepdom) = hold
    print "Final count:", total_count
    print "Err Count:", err_count
    print "total:", totalall, len(pdzdomains)
#import into database

if __name__ == '__main__':
    main()